# 自定义step_plot()
import numpy as np
import control as ctl
import matplotlib.pyplot as plt
from sympy import *
from numpy import linspace
from scipy.integrate import odeint, solve_bvp, solve_ivp
init_printing()


def step_plot(s):
    y, t = ctl.step_response(s)
    plt.plot(y, t, 'b', linewidth=1.5)
    plt.title('Step Response', fontsize=16)
    plt.xlabel('Time(seconds)', fontsize=12)
    plt.ylabel('Amplitude', fontsize=12)
    plt.show()


def impulse_plot(s):
    y, t = ctl.impulse_response(s)
    plt.plot(y, t, 'b', linewidth=1.5)
    plt.title('Impulse Response', fontsize=16)
    plt.xlabel('Time(seconds)', fontsize=12)
    plt.ylabel('Amplitude', fontsize=12)
    plt.show()


s = ctl.tf([4], [1, 2, 10, 8])
step_plot(s)
impulse_plot(s)


# 定义符号常量x 与 f(x)
x = Symbol('x')
f = symbols('f', cls=Function)
# 用diffeq代表微分方程： f''(x) + 3f'(x) + 2f(x) = exp(-x)
diffeq = Eq(f(x).diff(x, x) + 3*f(x).diff(x) + 2*f(x), exp(-x))
# 调用dsolve函数,返回一个Eq对象，hint控制精度
print(dsolve(diffeq, f(x)))
C1 = -2
C2 = 3


def fvdp1(t, y):
    dy1 = y[1]  # y[1]=dy/dt，一阶导
    dy2 = -3 * y[1] - 2 * y[0] + exp(-1 * t)
    # y[0]是最初始，也就是需要求解的函数
    # 注意返回的顺序是[一阶导， 二阶导]，这就形成了一阶微分方程组
    return [dy1, dy2]
# 或者下面写法更加简单


def fvdp2(t, y):
    y0, y1 = y
    # y0是需要求解的函数，y1是一阶导
    # 返回的顺序是[一阶导， 二阶导]，这就形成了一阶微分方程组
    dydt = [y1, -3 * y1 - 2 * y0 + exp(-t)]
    return dydt


def solve_second_order_ode():
    t2 = linspace(0, 20, 1000)
    tspan = (0, 20.0)
    y0 = [1.0, 2.0]  # 初值条件
    # 初值[2,0]表示y(0)=2,y'(0)=0
    # 返回y，其中y[:,0]是y[0]的值，就是最终解，y[:,1]是y'(x)的值
    y = odeint(fvdp1, y0, t2, tfirst=True)
    y_ = solve_ivp(fvdp2, t_span=tspan, y0=y0, t_eval=t2)

    plt.subplot(211)
    plt.title('Zero-state Response', fontsize=16)
    y1, = plt.plot(t2, y[:, 0], label='y')
    y1_1, = plt.plot(t2, y[:, 1], label='y‘')
    plt.legend(handles=[y1, y1_1])

    plt.subplot(212)
    y2, = plt.plot(y_.t, y_.y[0, :], 'g--', label='y(0)')
    y2_2, = plt.plot(y_.t, y_.y[1, :], 'r-', label='y(1)')
    plt.legend(handles=[y2, y2_2])
    plt.show()


solve_second_order_ode()
